# Successfully installed sentry_sdk-1.9.10
import json
import sentry_sdk
from fastapi import Request, Response, HTTPException, FastAPI
from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
from starlette.middleware.base import (BaseHTTPMiddleware, RequestResponseEndpoint)

from afast_core.core_plugins.pluginbase import IBasePlugin as BasePlugin


class SentryPluginClient(BasePlugin):
    # 设置插件默认的参数信息
    name = 'Sentry错误收集插件'

    def setup(self, app: FastAPI, name: str = None, *args, **kwargs):
        """插件初始化"""
        # 插件别名赋值到APP对象里面，让它持有该对象
        app.sentry = self
        # 插件关键参数校验
        assert 'sentry_setting' in kwargs or \
               'environment' in kwargs or \
               'release' in kwargs, '需要配置sentry_setting,environment,release等参数信息'
        # 业务逻辑初始化操作
        sentry_sdk.init(dsn=kwargs['sentry_setting'].SENTRY_DSN, integrations=[SqlalchemyIntegration()],
                        environment=kwargs.get('environment'), release=kwargs.get('release'))
        app.add_middleware(SentryMiddlware)


# 异步生成装饰器
class Aiwrap:
    def __init__(self, obj):
        self._it = iter(obj)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            value = next(self._it)
        except StopIteration:
            raise StopAsyncIteration
        return value


class SentryMiddlware(BaseHTTPMiddleware):

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        with sentry_sdk.configure_scope() as scope:
            try:
                response = await call_next(request)
                # 设置状态码
                status_code = response.__dict__['status_code']
                # 状态码
                if status_code >= 400 and status_code < 500:
                    # 获取到响应报文信息，但是必须使用基于BaseHTTPMiddleware的基础中间件
                    body_encode = [data async for data in response.__dict__['body_iterator']]
                    response.__setattr__('body_iterator', Aiwrap(body_encode))
                    body = json.loads(body_encode[0].decode())
                    scope.set_context('request', request)
                    scope.set_context('response', body)
                    scope.user = {'IP Address': request.client.host}
                    # 捕获异常
                    sentry_sdk.capture_exception(HTTPException(status_code=status_code, detail=body['detail']))

                return response

            except Exception as error:
                # 设置上下文对象
                scope.set_context('request', request)
                scope.user = {'IP Address': request.client.host}
                # 捕获异常
                sentry_sdk.capture_exception(error)
                raise error
